from abc import abstractmethod
from typing import List

from sam.util.dbUtil import DB, DEFAULT_MYSQL_CONNECTION_CONFIG, DbConnectionConfig
from sam.util.esUtil import EsConfig, EsUtil


class SyncDb2EsByGenerator(DB):
    def __init__(self, es_config: EsConfig, db_config: DbConnectionConfig = DEFAULT_MYSQL_CONNECTION_CONFIG):
        super().__init__(config=db_config)
        self.es = EsUtil(es_config)

    def get_db_data_by_generator(self
                                 , table_name
                                 , where: str = None
                                 , column: list = None
                                 , column_str: str = None
                                 , batch_size: int = 10000
                                 , optional: str = "list"
                                 , distinct: bool = False
                                 , datetime_formatter: str = None
                                 , result_formatter=None
                                 , is_only_print_sql: bool = False
                                 , is_fixed_start: bool = False
                                 ):
        if where and "status" in where:
            if not is_fixed_start:
                msg = "查询语句中使用了状态条件，但是没有查询固定开始位置"
                raise RuntimeWarning(msg)
        elif is_fixed_start:
            msg = "查询语句中固定了开始位置，但是查询条件中没有状态筛选条件"
            raise RuntimeWarning(msg)

        total_count = self.count(table_name, where=where, column=column, column_str=column_str, distinct=distinct)
        self.logger.info(f"从数据库中查询到 共有 {total_count} 条数据要处理")
        max_count = total_count - 1 + batch_size
        _start = 0
        _end = _start + batch_size
        while _end <= max_count:
            if is_fixed_start:
                __start = 0
            else:
                __start = _start
            yield self.query(
                table_name
                , where=where
                , column=column
                , column_str=column_str
                , start=__start
                , limit=batch_size
                , optional=optional
                , distinct=distinct
                , datetime_formatter=datetime_formatter
                , result_formatter=result_formatter
                , is_only_print_sql=is_only_print_sql
            )
            _start += batch_size
            _end += batch_size

    def batch_sync_db2_es_by_generator(self
                                       , table_name: str
                                       , where: str = None
                                       , column: list = None
                                       , column_str: str = None
                                       , batch_size: int = 50000
                                       , distinct: bool = False
                                       , datetime_formatter: str = None
                                       , is_fixed_start: bool = False
                                       , optional: str = "Update"
                                       , is_only_print_sql: bool = False
                                       , id_column: str = "id"
                                       ):
        loop = self.get_db_data_by_generator(
            table_name, where=where, column=column, column_str=column_str, batch_size=batch_size,
            distinct=distinct, optional="dict", is_fixed_start=is_fixed_start
            , datetime_formatter=datetime_formatter
        )
        while True:
            try:
                one_generator_return_list = next(loop)
                self._handle_one_generator_return_data(one_generator_return_list, optional=optional,
                                                       is_only_print_sql=is_only_print_sql, id_column=id_column)
            except StopIteration:
                break

    @abstractmethod
    def pre_handle_db_query_result_list(self, pre_result_list: List[dict]) -> List[dict]:
        pass

    def _handle_one_generator_return_data(self, one_generator_return_list, optional: str = "Update",
                                          is_only_print_sql: bool = False
                                          , id_column: str = "id"
                                          , batch_size: int = 100):
        result_list_init = self.pre_handle_db_query_result_list(one_generator_return_list)
        result_list = [item for item in result_list_init if item]
        if "Update" == optional:
            self.es.batch_update_es(
                result_list
                , is_only_print_sql=is_only_print_sql
                , batch_size=batch_size
                , id_column=id_column
            )
        elif "Delete" == optional:
            self.es.batch_delete_es(
                result_list
                , is_only_print_sql=is_only_print_sql
                , batch_size=batch_size
                , id_column=id_column
            )
        elif "Insert" == optional:
            self.es.batch_insert_or_update_es(
                result_list
                , is_only_print_sql=is_only_print_sql
                , batch_size=batch_size
                , id_column=id_column
            )
        else:
            msg = f"只支持[Update, Delete, Insert]操作, 不支持当前操作{optional}"
            raise RuntimeError(msg)
